
package com.chungkui.ratelimiter;

import com.chungkui.check.util.DigestUtils;
import org.springframework.core.io.ClassPathResource;
import org.springframework.scripting.support.ResourceScriptSource;

import java.io.IOException;
import java.time.Instant;
import java.util.Arrays;
import java.util.List;
import java.util.Objects;


/**
 * 基于redis流控.<br>
 *
 * @author jason
 * @see [相关类/方法](可选)
 * @since [产品/模块版本] (可选)
 */
public abstract class AbstractRedisRateLimiterImpl implements RedisRateLimiter {
    private static volatile String redisScriptSha;

    @Override
    public boolean check(String key, int rate, int capacity, int perUser) throws IOException {
        if (key == null || rate == 0) {
            return true;
        }
        if (capacity == 0) {
            capacity = rate;
        }
        if (perUser == 0) {
            perUser = 1;
        }
        List<Long> res;
        List<String> keys = getKeys(key);
        if (redisScriptSha == null) {
            synchronized (this) {
                String src = new ResourceScriptSource(
                        new ClassPathResource("META-INF/scripts/chungkui_rate_limiter.lua")).getScriptAsString();
                res = eval(src, 2, keys.get(0), keys.get(1), rate + "", capacity + "", Instant.now().getEpochSecond() + "", perUser + "");
                redisScriptSha = DigestUtils.sha1DigestAsHex(src);
                return Objects.equals(res.get(0), 1L);
            }
        }
        res = evalSha(redisScriptSha, 2, keys.get(0), keys.get(1), rate + "", capacity + "", Instant.now().getEpochSecond() + "", perUser + "");
        return Objects.equals(res.get(0), 1L);
    }

    static List<String> getKeys(String id) {
        String prefix = "chungkui_rate_limiter.{" + id;
        String tokenKey = prefix + "}.tokens";
        String timestampKey = prefix + "}.timestamp";
        return Arrays.asList(tokenKey, timestampKey);
    }

    /**
     * redis 执行入口。需要你自行实现
     *
     * @param redisScript lua脚本
     * @param keySize
     * @param scriptArgs
     * @return
     */
    public abstract List<Long> eval(String redisScript, int keySize, String... scriptArgs);

    /**
     * @param sha
     * @param keySize
     * @param scriptArgs
     * @return
     */
    public abstract List<Long> evalSha(String sha, int keySize, String... scriptArgs);


}
